/* -- INCLUDES ---------------------------------------------------------------*/
#include <stdio.h>
#include <sys/types.h>
#include <unistd.h>
#include <stdlib.h>
#include <signal.h>
#include <getopt.h>
#include <ctype.h>
#include <string.h>
#include <sys/msg.h>
#include <sys/wait.h>
#include <errno.h>

#include "Fuzzit.h"
#include "Watcher.h"
#include "Logger.h"
#include "Launcher.h"
#include "Utils.h"

/* -- DEFINES ----------------------------------------------------------------*/

/* -- PRIVATE VARIABLES ------------------------------------------------------*/
static int msgq_logs			= -1;
static int msgq_files			= -1;
static char * input_dir			= NULL;
static char * executable		= NULL;
static char * log_file			= NULL;
static bool_t fuzz_with_content	= false;
static unsigned int max_time	= 0;
static unsigned int arg_index	= 0;
static char ** args				= {NULL, NULL};
static char ** envp				= {NULL};

/* -- PRIVATE FUNCTIONS ------------------------------------------------------*/
static void usage( char * prog )
{
	printf("Usage : %s [OPTIONS]\n", prog ? prog : DEFAULT_NAME );
	printf("\t-d, --dir <dir>\tdirectory to monitor for files (default is %s)\n", DEFAULT_FUZZ_DIR);
	printf("\t-e, --exec <file>\texecutable to fuzz (required)\n");
	printf("\t-l, --log <file>\tlog file for the fuzzing session (default is %s)\n", DEFAULT_LOG_FILE);
	printf("\t-t, --time <value>\tmaximum execution time (in sec) for each fuzzing process (default is %d sec)\n", DEFAULT_MAXTIME);
	printf("\t-a, --argv [0:arg0 1:arg2 ... X:argX]\tcommand line arguments\n", DEFAULT_MAXTIME);
	printf("\t-n, --env [0:env0 1:env1 ... X:envX]\tenvironnement variables\n");
	printf("\t-c, --content\tfuzzing with file content instead of file names\n");
	printf("\t-h, --help\tthis help\n");

	exit(EXIT_FAILURE);
}

static char ** parse_array( char * str )
{
}

static void handle_options( int argc, char ** argv )
{
	struct option options[] = {
		{"dir",		required_argument,	NULL,	'd'},
		{"exec",	required_argument,	NULL,	'e'},
		{"log",		required_argument,	NULL,	'l'},
		{"time",	required_argument,	NULL,	't'},
		{"args",	required_argument,	NULL,	'a'},
		{"env",		required_argument,	NULL,	'n'},
		{"content",	no_argument,		NULL,	'c'},
		{"help",	no_argument,		NULL,	'h'},
		{0,			0,					0,		0}
	};

	char c;
	int index;

	while( (c = getopt_long(argc, argv, "d:e:l:t:a:n:ch", options, &index)) != -1 ) {
		switch( c ) {
			case 'd':
				input_dir = strdup(optarg);
				break;

			case 'e':
				executable = strdup(optarg);
				break;
				
			case 'l':
				log_file = strdup(optarg);
				break;

			case 't':
				max_time = atoi(optarg);
				break;

			case 'a':
				args = parse_array(optarg);
				break;

			case 'n':
				envp = parse_array(optarg);
				break;
				
			case 'c':
				fuzz_with_content = true;
				break;
				
			case 'h':
				usage( NULL );
				break;

			default	:
				if( isalnum(c) )
					FFATAL("[-] Unknow option \"%c\"\n", c);
				else
					FATAL("[-] Unknow option\n");
		}
	}

	if( input_dir == NULL )
		input_dir = strdup(DEFAULT_FUZZ_DIR);

	if( log_file == NULL )
		log_file = strdup(DEFAULT_LOG_FILE);
		
	if( max_time <= 0 )
		max_time = DEFAULT_MAXTIME;

	if( executable == NULL )
		FATAL("You must specify an exectable with option -e\n");

	/* TODO : handle others options */
}

static void create_ressources()
{
	msgq_logs = msgget(MSGQ_LOGS_KEY, IPC_CREAT | IPC_EXCL | IPC_RIGHTS);
	msgq_files = msgget(MSGQ_FILES_KEY, IPC_CREAT | IPC_EXCL | IPC_RIGHTS);

	if( msgq_logs == -1 || msgq_files == -1 )
		PFATAL("[-] Cannot create message queues ");

	DEBUG("[~] msgq_logs=%u, msgq_files=%u\n", msgq_logs, msgq_files);

	/* extend max size of queues ? */
}

static void destroy_ressources()
{
	msgctl(msgq_logs, IPC_RMID, NULL);
	msgctl(msgq_files, IPC_RMID, NULL);

	if( input_dir != NULL )
		free( input_dir );
}

static rlim_t extend_limit( int ressource, char * description )
{
	struct rlimit lim;

	if( getrlimit( ressource, &lim ) != 0 ) {
		ERROR( "[-] Cannot retreive limit for %s : %s\n", description, strerror(errno) );
		return 0;
	} else if( lim.rlim_cur == RLIM_INFINITY ) {
		printf( "[+] %s limit is already RLIM_INFINITY\n", description );
		return RLIM_INFINITY;
	} else {
		lim.rlim_cur = lim.rlim_max;
		if( setrlimit( ressource, &lim ) != 0 ) {
			ERROR( "[-] Cannot extend %s limit\n", description );
			return lim.rlim_max;
		} else if( lim.rlim_max == RLIM_INFINITY ) {
			printf( "[+] %s limit extended to RLIM_INFINITY\n", description );
			return RLIM_INFINITY;
		} else {
			printf( "[+] %s limit extended to %u\n", description, (unsigned int)lim.rlim_max );
			return lim.rlim_max;
		}
	}
}

static void adjust_priority( rlim_t nice_limit )
{
	int inc;

	if( nice_limit != 0 ) {
		if( nice_limit == RLIM_INFINITY )
			inc = -20;
		else
			inc = 20 - nice_limit;

		if( nice(inc) == -1 )
			ERROR( "[-] Cannot adjust priority (%d) : %s\n", inc, strerror(errno) );
		else
			printf( "[+] Priority adjusted (%d)\n", inc );
	}
}

static void init_core_pattern()
{
	FILE * pattern_file = fopen(CORE_PATTERN_FILE, "w");
	
	if( pattern_file == NULL )
		ERROR( "[-] Cannot open core pattern file : %s\n", strerror(errno) );
	else if( fputs(CORE_PATTERN, pattern_file) == EOF )
		ERROR( "[-] Cannot write core pattern : %s\n", strerror(errno) );
	else
		printf( "[+] Core pattern written : %s\n", CORE_PATTERN );
}

/* -- PUBLIC FUNCTIONS -------------------------------------------------------*/
int main( int argc, char ** argv )
{
	rlim_t nice_limit;
	pid_t logger = -1, watcher = -1, launcher = -1;
	
	if( argc < COUNT_REQUIRED_ARGS+1 )
		usage( argv[0] );

	handle_options(argc, argv);
	create_ressources();
	
	extend_limit( RLIMIT_CORE,	"Core size" );
	extend_limit( RLIMIT_NPROC,	"Process/Tread" );
	nice_limit = extend_limit( RLIMIT_NICE,	"Nice" );
	
	adjust_priority( nice_limit );
	init_core_pattern();

	printf("[+] Fuzzing executable \"%s\"\n", executable);
	printf("[+] Monitoring directory \"%s\"\n", input_dir);
	printf("[+] Log file \"%s\"\n", log_file);

	logger = fork();
	switch( logger ) {
		case -1	:
			PERROR("[-] Failed to fork logger");
			goto end;

		case 0	:
			Logger( msgq_logs, log_file );
			PFATAL("[-] Logger returned ");

		default	:
			launcher = fork();
			switch( launcher ) {
				case -1	:
					PERROR("[-] Failed to fork launcher");
					goto end;

				case 0	:
					Launcher( msgq_files, executable, args, envp, arg_index, fuzz_with_content, max_time );
					PFATAL("[-] Launcher returned ");

				default	:
					watcher = fork();
					switch( watcher ) {
						case -1	:
							PERROR("[-] Failed to fork watcher");
							goto end;

						case 0	:
							Watcher( input_dir, msgq_files );
							PFATAL("[-] Watcher returned ");

						default	:
							DEBUG("[~] logger=%d, launcher=%d, watcher=%d\n", logger, launcher, watcher);
							printf("[+] Initialisation done. Now sleeping, \"fg\" to resume & stop fuzzing\n");
							kill( getpid(), SIGSTOP );
							/* Sleep 'til I wake you up bitch ! */
					}
			}
	}

end:
	printf("[+] Done. Now exiting.\n");

	KILLWAIT( watcher, SIGTERM );
	KILLWAIT( launcher, SIGTERM );
	KILLWAIT( logger, SIGTERM );

	destroy_ressources();

	return EXIT_SUCCESS;
}
